import json

import requests
from more_itertools import chunked
from tqdm import tqdm

from iflytech_assistant.assistant.heqc import load_tags_from_excel

preference_tags, refined_tags = load_tags_from_excel(
    ".vscode/高情商沟通/生成偏好+修改词-家人改长辈.xlsx"
)
dataset_id = "58824550-4382-4f3d-9c18-d14a761dd4bd"
document_id = "45f13257-55b4-4c1f-8e5b-81f7523acc9e"

url = f"http://106.52.220.44/:5001/v1/datasets/{dataset_id}/documents/{document_id}/segments"

headers = {
    "Authorization": "Bearer dataset-m8azj29Wod7dGiV5qn6AIud7",
    "Content-Type": "application/json",
}

for target, tags in preference_tags.items():
    for batch in tqdm(chunked(tags, 5), total=len(tags) // 5, desc=target):
        data = {
            "segments": [
                {
                    "content": f"{target} {tag.tag}",
                    "answer": json.dumps(tag.model_dump(), ensure_ascii=False),
                    "keywords": [target, tag.tag],
                }
                for tag in batch
            ]
        }

        response = requests.post(
            url, headers=headers, data=json.dumps(data, ensure_ascii=False)
        )
